/*
** mytftp-client.c
** by Matt Rucker
** Created: Mar 17, 2013
** Last modified: May 6th, 2013
*/

#include<stdio.h>
#include<stdlib.h>
#include<unistd.h>
#include<errno.h>
#include<string.h>
#include<sys/types.h>
#include<sys/socket.h>
#include<netinet/in.h>
#include<arpa/inet.h>
#include<netdb.h>

#define MAXDATALEN 512
#define DEFAULTPORT 3335

#define RRQ 1
#define WRQ 2
#define DATA 3
#define ACK 4
#define ERROR 5
#define NET 0
#define OCT 1
#define MAIL 2
#define INITIAL_SIZE 4

struct Buffer {
    char* data;
    int next;
    size_t size;
};

void reserve_space(struct Buffer* b, size_t bytes);
void serialize_short(short x, struct Buffer* b);
void serialize_str(char* x, struct Buffer* b);
void serialize_char(char x, struct Buffer* b);
struct Buffer* new_buffer();

int main( int argc, char * argv[] ) {

    char* modes[] = {"netascii", "octet", "mail"};
    char* mode;
    char* filename;
    int sockfd ;
    int their_port; 
    struct sockaddr_in my_addr; 
    struct sockaddr_in their_addr ; 
    struct hostent *he ; 
    struct Buffer* buf = new_buffer();
    int numbytes ;
    int addr_len; 
    int i;
    unsigned short flags[4] = {0}; 
    FILE *file;
    short block = 0;
    short opcode;
    short rq; 
    long c;
	int count = 1;

    struct timeval tv;
    tv.tv_sec = 5;
    tv.tv_usec = 0;

	if ((sockfd = socket(AF_INET, SOCK_DGRAM, 0)) == -1 ) {
		perror("socket") ;
		exit(1) ;
	}

    /*if (setsockopt(sockfd, SOL_SOCKET, SO_RCVTIMEO,&tv,sizeof(tv)) < 0) {
        perror("Error");
    }*/

    for(i=1;i<argc;i++) {
        if(!strcmp(argv[i],"-l")) flags[0] = 1;
        else if(!strcmp(argv[i],"-p")) flags[1] = atoi(argv[i+1]);
        else if(!strcmp(argv[i],"-v")) flags[2] = 1;
        else if(!strcmp(argv[i],"-r")) flags[3] = RRQ;
        else if(!strcmp(argv[i],"-w")) flags[3] = WRQ;
    }
    if(flags[0]) {

        my_addr.sin_family = AF_INET ;
        my_addr.sin_port = (flags[1]) ? htons(flags[1]) : htons(DEFAULTPORT);
        my_addr.sin_addr.s_addr = INADDR_ANY ;
        memset(&(my_addr.sin_zero),'\0',8) ;

        if (bind(sockfd, (struct sockaddr *)&my_addr, 
	          sizeof(struct sockaddr)) == -1 ) {
	        perror("bind") ;
	        exit(1) ;
	    }
        reserve_space(buf,MAXDATALEN*sizeof(char));

            addr_len = sizeof(struct sockaddr) ;
	    if ((numbytes=recvfrom(sockfd, buf->data, buf->size,0,
		  (struct sockaddr *)&their_addr, &addr_len)) == -1 ) {
            perror("recvfrom") ;
		    exit(1) ;
            }

        memcpy(&opcode,buf->data,sizeof(short));
        opcode = ntohs(opcode);
        buf->next = 2;
        while(buf->next < numbytes) {
            memcpy(&c,buf->data + buf->next,sizeof(char));
            if(c=='\0'){
                mode = (char*)malloc(sizeof(char)*(buf->next-sizeof(short)+1));
                memcpy(mode,buf->data + sizeof(short),buf->next-sizeof(short));

                filename = (char*)malloc(sizeof(char)*(numbytes-buf->next));
                memcpy(filename,buf->data + buf->next+1,9);

                break;
            }
            buf->next++;
        }

        free(buf);
        buf = new_buffer();
        if(opcode==RRQ){
            if((file=fopen(filename,"rb"))==NULL) perror("fopen");
	        while(i!=-1){
				block++;
				serialize_short(DATA,buf);
				serialize_short(block,buf);
				i = 0;
				while(i < MAXDATALEN){
					if((c=fgetc(file))==EOF) {i=-1;break;}
					serialize_char(c,buf);
					i++;
				}
				if ((numbytes=sendto(sockfd, buf->data, buf->next,0,
				  (struct sockaddr *)&their_addr, sizeof(struct sockaddr)) ) == -1 ) {
					perror("sendto") ;
					exit(1) ;
				}
				free(buf);
				buf = new_buffer();
			    addr_len = sizeof(struct sockaddr) ;
		   	    if ((numbytes=recvfrom(sockfd, buf->data, buf->size,0,
				  (struct sockaddr *)&their_addr, &addr_len)) == -1 ) {
			    	perror("recvfrom") ;
					exit(1) ;
			    }
                memcpy(&opcode,buf->data,sizeof(short));
                memcpy(&block,buf->data + sizeof(short),sizeof(short));
                opcode = ntohs(opcode);
                block = ntohs(block);
				free(buf);
				buf = new_buffer();	
            }
        }
        else if(opcode==WRQ){
            if((file=fopen(filename,"wb"))==NULL) perror("fopen");
   			serialize_short(ACK,buf);
    		serialize_short(block,buf);
		    if ((numbytes=sendto(sockfd, buf->data, buf->next,0,
			  (struct sockaddr *)&their_addr, sizeof(struct sockaddr)) ) == -1 ) {
				perror("sendto") ;
				exit(1) ;
		    }
            while(i!=-1){
				free(buf);
				buf = new_buffer();
                reserve_space(buf,MAXDATALEN+(sizeof(short)*2));
                addr_len = sizeof(struct sockaddr) ;
	            if ((numbytes=recvfrom(sockfd, buf->data, buf->size,0,
                  (struct sockaddr *)&their_addr, &addr_len)) == -1 ) {
                    perror("recvfrom") ;
		            exit(1) ;
                }
				if (numbytes<(MAXDATALEN+(sizeof(char)*2))) {i=-1;}
                memcpy(&opcode,buf->data,sizeof(short));
                memcpy(&block,buf->data + sizeof(short),sizeof(short));
                opcode = ntohs(opcode);
                block = ntohs(block);
                if(opcode==DATA) {
                    buf->next = 4;
                    while(buf->next < numbytes){
                        memcpy(&c,buf->data + buf->next,sizeof(char));
                        buf->next++;
                        fputc(c,file);
                    }
                }
                free(buf);
                buf = new_buffer();
				block++;
                serialize_short(ACK,buf);
                serialize_short(block,buf);
                if ((numbytes=sendto(sockfd, buf->data, buf->next,0,
		          (struct sockaddr *)&their_addr, sizeof(struct sockaddr)) ) == -1 ) {
	                perror("sendto") ;
                    exit(1) ;
                }
            }
        } else if (opcode==ERROR) printf("Error!\n");
    } else {
		i = 0;
        rq = flags[3];
        filename = (char*)malloc((strlen(argv[argc-2])+1)*sizeof(char));
        strcpy(filename,argv[argc-2]);        

        if ((he=gethostbyname(argv[argc-1]))==NULL) {
	        perror("gethostbyname") ;
		    exit(1) ;
	    }

        their_addr.sin_family = AF_INET ;
        their_addr.sin_port = (flags[1]) ? htons(flags[1]) : htons(DEFAULTPORT);
        their_addr.sin_addr = *((struct in_addr *)he->h_addr) ;
        memset(&(their_addr.sin_zero), '\0', 8 ) ;

        serialize_short(rq,buf);
        serialize_str(modes[OCT],buf);
        serialize_str(filename,buf);
        if ((numbytes=sendto(sockfd, buf->data, buf->next,0,
		  (struct sockaddr *)&their_addr, sizeof(struct sockaddr)) ) == -1 ) {
		    perror("sendto") ;
		    exit(1) ;
        }

        free(buf);
        buf = new_buffer();
        if(rq==RRQ){
	        if((file=fopen(argv[argc-2],"wb"))==NULL)
	            perror("fopen");
            while(i!=-1){
				reserve_space(buf,MAXDATALEN+(sizeof(short)*2));
				addr_len = sizeof(struct sockaddr);
				if ((numbytes=recvfrom(sockfd, buf->data, buf->size,0,
					  (struct sockaddr *)&their_addr, &addr_len)) == -1 ) {
					perror("recvfrom") ;
						exit(1) ;
				}
				memcpy(&opcode,buf->data,sizeof(short));
				memcpy(&block,buf->data + sizeof(short),sizeof(short));
				opcode = ntohs(opcode);
				block = ntohs(block);

				if(opcode==DATA && block==count) {
					for(i=4;i<numbytes;i++) fputc(*(buf->data + i),file);

				    if (numbytes<MAXDATALEN+(sizeof(short)*2)) i=-1;
					free(buf);
				    buf = new_buffer();
		            
		   			serialize_short(ACK,buf);
		    		serialize_short(block,buf);
					if ((numbytes=sendto(sockfd, buf->data, buf->next,0,
			  		  (struct sockaddr *)&their_addr, sizeof(struct sockaddr)) ) == -1 ) {
						perror("sendto") ;
						exit(1) ;
		    		}
					count++;
				} else if (opcode==ERROR) {printf("Error\n");}
            }
        } else if(rq==WRQ) {
			if((file=fopen(argv[argc-2],"rb"))==NULL)
	            perror("fopen");
			addr_len = sizeof(struct sockaddr);
			if ((numbytes=recvfrom(sockfd, buf->data, buf->size,0,
				  (struct sockaddr *)&their_addr, &addr_len)) == -1 ) {
				perror("recvfrom") ;
					exit(1) ;
			}
			memcpy(&opcode,buf->data,sizeof(short));
			memcpy(&block,buf->data + sizeof(short),sizeof(short));
			opcode = ntohs(opcode);
			block = ntohs(block);
			free(buf);
			buf = new_buffer();
			if (opcode==ACK) {
				block++;
				while(i!=-1){
					reserve_space(buf,MAXDATALEN+(sizeof(short)*2));
		   			serialize_short(DATA,buf);
		    		serialize_short(block,buf);
					i=0;
					while(i < MAXDATALEN){
						if((c=fgetc(file))==EOF) {i=-1;break;}
						serialize_char(c,buf);
						i++;
					}
					if ((numbytes=sendto(sockfd, buf->data, buf->next,0,
			  		  (struct sockaddr *)&their_addr, sizeof(struct sockaddr)) ) == -1 ) {
						perror("sendto") ;
						exit(1) ;
		    		}
					free(buf);
					buf = new_buffer();
					addr_len = sizeof(struct sockaddr);
					if ((numbytes=recvfrom(sockfd, buf->data, buf->size,0,
						  (struct sockaddr *)&their_addr, &addr_len)) == -1 ) {
						perror("recvfrom") ;
						exit(1) ;
					}
					memcpy(&opcode,buf->data,sizeof(short));
					memcpy(&block,buf->data + sizeof(short),sizeof(short));
					opcode = ntohs(opcode);
					block = ntohs(block);
					free(buf);
					buf = new_buffer();
				}
			} else if (opcode==ERROR) {printf("Error\n");}
		}
    }
    free(buf);
    fclose(file);
    close(sockfd) ;
	return 0 ;
}

void reserve_space(struct Buffer* b, size_t bytes) {
	while((b->next + bytes) > b->size) {
		b->data = realloc(b->data,b->size * 2);
        b->size *= 2;
	}
}

void serialize_short(short x, struct Buffer* b) {
    x = htons(x);

    reserve_space(b,sizeof(short));
    memcpy(((char *)b->data) + b->next, &x, sizeof(short));
    b->next += sizeof(short);
}

void serialize_str(char* x, struct Buffer* b) {
    reserve_space(b,strlen(x)+1);
    memcpy(((char *)b->data) + b->next,x,strlen(x)+1);
    b->next += (int)strlen(x) + 1;
}

void serialize_char(char x, struct Buffer* b) {
    reserve_space(b,1);
    memcpy(((char *)b->data) + b->next,&x,1);
    b->next += 1;
}

struct Buffer* new_buffer() {
    struct Buffer *b = malloc(sizeof(struct Buffer));
 
    b->data = malloc(INITIAL_SIZE);
    b->size = INITIAL_SIZE;
    b->next = 0;

    return b;
};
